package com.simple.mqtt.server.processor;

import com.simple.mqtt.common.domain.TopicFilter;
import com.simple.mqtt.common.task.MqttPendingPublish;
import com.simple.mqtt.common.task.MqttPendingQos2Publish;
import com.simple.mqtt.common.tool.LogTool;
import com.simple.mqtt.server.config.MqttBrokerConfig;
import com.simple.mqtt.server.domain.Message;
import com.simple.mqtt.server.domain.MessageType;
import com.simple.mqtt.server.domain.Subscribe;
import com.simple.mqtt.server.serializer.IMessageSerializer;
import io.netty.channel.ChannelHandlerContext;
import io.netty.handler.codec.mqtt.*;
import io.netty.util.AttributeKey;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ScheduledExecutorService;

/**
 * MQTT协议处理器
 * 处理客户端的所有MQTT消息类型，包括CONNECT、PUBLISH、SUBSCRIBE、UNSUBSCRIBE、PINGREQ、DISCONNECT
 * 管理消息确认机制（PUBACK/PUBREC/PUBREL/PUBCOMP）
 */
public class MqttBrokerProcessor {
    /** 存储在ChannelHandlerContext中的客户端ID属性键 */
    public static final AttributeKey<String> CLIENT_ID_ATTR = AttributeKey.valueOf("clientId");
    /** 存储在ChannelHandlerContext中的用户名属性键 */
    public static final AttributeKey<String> USERNAME_ATTR = AttributeKey.valueOf("username");

    private final MqttBrokerConfig config;
    private final ScheduledExecutorService scheduledExecutor;

    /**
     * 构造函数
     * @param config MQTT代理配置
     * @param scheduledExecutor 调度执行器，用于消息重传和异步任务
     */
    public MqttBrokerProcessor(MqttBrokerConfig config, ScheduledExecutorService scheduledExecutor) {
        this.config = config;
        this.scheduledExecutor = scheduledExecutor;
    }

    /**
     * 处理CONNECT消息
     * @param ctx 通道上下文
     * @param msg 连接消息
     */
    public void processConnect(ChannelHandlerContext ctx, MqttConnectMessage msg) {
        MqttConnectPayload payload = msg.payload();
        String clientId = payload.clientIdentifier();
        String username = payload.userName();
        String password = payload.passwordInBytes() != null ?
                new String(payload.passwordInBytes()) : null;

        if (clientId == null || clientId.isEmpty()) {
            sendConnAck(ctx, MqttConnectReturnCode.CONNECTION_REFUSED_IDENTIFIER_REJECTED);
            return;
        }

        if (clientId.length() > config.getMaxClientIdLength()) {
            sendConnAck(ctx, MqttConnectReturnCode.CONNECTION_REFUSED_IDENTIFIER_REJECTED);
            return;
        }

        if (!config.getAuthHandler().authenticate(username, password)) {
            sendConnAck(ctx, MqttConnectReturnCode.CONNECTION_REFUSED_BAD_USER_NAME_OR_PASSWORD);
            return;
        }

        ChannelHandlerContext existingCtx = config.getChannelManager().getChannel(clientId);
        if (existingCtx != null && existingCtx != ctx) {
            LogTool.warn("Client {} already connected, closing old connection", clientId);
            existingCtx.close();
        }

        ctx.channel().attr(CLIENT_ID_ATTR).set(clientId);
        if (username != null) {
            ctx.channel().attr(USERNAME_ATTR).set(username);
        }
        config.getChannelManager().addChannel(clientId, ctx);

        if (msg.variableHeader().isWillFlag()) {
            Message willMessage = new Message();
            willMessage.setClientId(clientId);
            willMessage.setTopic(payload.willTopic());
            willMessage.setPayload(payload.willMessageInBytes());
            willMessage.setQos(msg.variableHeader().willQos());
            willMessage.setRetain(msg.variableHeader().isWillRetain());
            willMessage.setMessageType(MessageType.DOWN_STREAM);
            config.getMessageStore().addWillMessage(clientId, willMessage);
        }

        MqttConnAckMessage connAck = MqttMessageBuilders.connAck()
                .returnCode(MqttConnectReturnCode.CONNECTION_ACCEPTED)
                .sessionPresent(false)
                .build();
        ctx.writeAndFlush(connAck);

        LogTool.info("Client {} connected, username: {}", clientId, username);
    }

    /**
     * 处理PUBLISH消息
     * @param ctx 通道上下文
     * @param msg 发布消息
     * @return 如果处理成功返回true
     */
    public boolean processPublish(ChannelHandlerContext ctx, MqttPublishMessage msg) {
        try {
            String clientId = ctx.channel().attr(CLIENT_ID_ATTR).get();
            if (clientId == null) {
                ctx.close();
                return false;
            }

            String topic = msg.variableHeader().topicName();
            MqttQoS qos = msg.fixedHeader().qosLevel();
            boolean retain = msg.fixedHeader().isRetain();
            byte[] payload = new byte[msg.payload().readableBytes()];
            msg.payload().readBytes(payload);

            if (retain) {
                Message retainMessage = new Message();
                retainMessage.setTopic(topic);
                retainMessage.setPayload(payload);
                retainMessage.setQos(qos.value());
                retainMessage.setRetain(false);
                retainMessage.setMessageType(MessageType.DOWN_STREAM);
                config.getMessageStore().addRetainMessage(topic, 0, retainMessage);
            }

            int messageId = msg.variableHeader().packetId();
            if (qos == MqttQoS.AT_LEAST_ONCE && messageId > 0) {
                MqttFixedHeader pubAckFixedHeader = new MqttFixedHeader(MqttMessageType.PUBACK, false, MqttQoS.AT_MOST_ONCE, false, 0);
                MqttMessageIdVariableHeader pubAckVariableHeader = MqttMessageIdVariableHeader.from(messageId);
                MqttMessage pubAck = new MqttMessage(pubAckFixedHeader, pubAckVariableHeader);
                ctx.writeAndFlush(pubAck);
            } else if (qos == MqttQoS.EXACTLY_ONCE && messageId > 0) {
                MqttFixedHeader pubRecFixedHeader = new MqttFixedHeader(MqttMessageType.PUBREC, false, MqttQoS.AT_MOST_ONCE, false, 0);
                MqttMessage pubRec = new MqttMessage(pubRecFixedHeader, MqttMessageIdVariableHeader.from(messageId));

                MqttPendingQos2Publish pendingQos2Publish = new MqttPendingQos2Publish(msg, pubRec);
                config.getSessionManager().addPendingQos2Publish(clientId, messageId, pendingQos2Publish);
                pendingQos2Publish.startPubRecRetransmitTimer(scheduledExecutor, ctx);

                ctx.writeAndFlush(pubRec);
            }

            // 使用消息分发器发布给订阅者
            Message message = new Message();
            message.setTopic(topic);
            message.setPayload(payload);
            message.setQos(qos.value());
            message.setRetain(retain);
            message.setMessageType(MessageType.DOWN_STREAM);
            config.getMessageDispatcher().dispatch(message, clientId);
            
            return true;
        } catch (Exception e) {
            LogTool.error("Error processing PUBLISH message", e);
            return false;
        }
    }

    /**
     * 处理PUBACK消息（QoS 1确认）
     * @param ctx 通道上下文
     * @param variableHeader 消息ID头
     */
    public void processPubAck(ChannelHandlerContext ctx, MqttMessageIdVariableHeader variableHeader) {
        String clientId = ctx.channel().attr(CLIENT_ID_ATTR).get();
        if (clientId == null) {
            return;
        }

        int messageId = variableHeader.messageId();
        MqttPendingPublish pendingPublish = config.getSessionManager().getPendingPublish(clientId, messageId);
        if (pendingPublish != null) {
            pendingPublish.onPubAckReceived();
            config.getSessionManager().removePendingPublish(clientId, messageId);
        }
    }

    /**
     * 处理PUBREC消息（QoS 2接收确认）
     * @param ctx 通道上下文
     * @param variableHeader 消息ID头
     */
    public void processPubRec(ChannelHandlerContext ctx, MqttMessageIdVariableHeader variableHeader) {
        String clientId = ctx.channel().attr(CLIENT_ID_ATTR).get();
        if (clientId == null) {
            return;
        }

        int messageId = variableHeader.messageId();
        MqttPendingPublish pendingPublish = config.getSessionManager().getPendingPublish(clientId, messageId);
        if (pendingPublish != null) {
            // Stop publish message retransmission timer
            pendingPublish.stopRetransmission();

            // Create PubRel message and start retransmission timer
            MqttFixedHeader fixedHeader = new MqttFixedHeader(MqttMessageType.PUBREL, false, MqttQoS.AT_LEAST_ONCE, false, 0);
            MqttMessage pubRelMessage = new MqttMessage(fixedHeader, variableHeader);

            pendingPublish.setPubRelMessage(pubRelMessage);
            pendingPublish.startPubRelRetransmissionTimer(scheduledExecutor, ctx);
            ctx.writeAndFlush(pubRelMessage);
        }
    }

    /**
     * Process PUBREL message (QoS 2 release)
     * @param ctx channel context
     * @param variableHeader message ID header
     */
    public void processPubRel(ChannelHandlerContext ctx, MqttMessageIdVariableHeader variableHeader) {
        String clientId = ctx.channel().attr(CLIENT_ID_ATTR).get();
        if (clientId == null) {
            return;
        }

        int messageId = variableHeader.messageId();
        MqttPendingQos2Publish pendingQos2Publish = config.getSessionManager().getPendingQos2Publish(clientId, messageId);
        if (pendingQos2Publish != null) {
            pendingQos2Publish.onPubRelReceived();
            config.getSessionManager().removePendingQos2Publish(clientId, messageId);
        }

        // 创建PubComp消息
        MqttFixedHeader fixedHeader = new MqttFixedHeader(MqttMessageType.PUBCOMP, false, MqttQoS.AT_MOST_ONCE, false, 0);
        MqttMessage pubCompMessage = new MqttMessage(fixedHeader, variableHeader);
        ctx.writeAndFlush(pubCompMessage);
    }

    public void processPubComp(ChannelHandlerContext ctx, MqttMessageIdVariableHeader variableHeader) {
        String clientId = ctx.channel().attr(CLIENT_ID_ATTR).get();
        if (clientId == null) {
            return;
        }

        int messageId = variableHeader.messageId();
        MqttPendingPublish pendingPublish = config.getSessionManager().getPendingPublish(clientId, messageId);
        if (pendingPublish != null) {
            pendingPublish.onPubCompReceived();
            config.getSessionManager().removePendingPublish(clientId, messageId);
        }
    }

    public void processSubscribe(ChannelHandlerContext ctx, MqttSubscribeMessage msg) {
        String clientId = ctx.channel().attr(CLIENT_ID_ATTR).get();
        if (clientId == null) {
            ctx.close();
            return;
        }

        List<MqttTopicSubscription> topicSubscriptions = msg.payload().topicSubscriptions();
        List<Integer> grantedQosLevels = new ArrayList<>();

        for (MqttTopicSubscription subscription : topicSubscriptions) {
            String topicFilter = subscription.topicName();
            MqttQoS requestedQos = subscription.qualityOfService();

            TopicFilter filter = new TopicFilter(topicFilter);
            config.getSessionManager().addSubscribe(filter, clientId, requestedQos.value());
            grantedQosLevels.add(requestedQos.value());

            LogTool.debug("Client {} subscribed to topic {} with QoS {}", clientId, topicFilter, requestedQos);
        }

        MqttFixedHeader subAckFixedHeader = new MqttFixedHeader(MqttMessageType.SUBACK, false, MqttQoS.AT_MOST_ONCE, false, 0);
        MqttMessageIdVariableHeader subAckVariableHeader = MqttMessageIdVariableHeader.from(msg.variableHeader().messageId());
        // Create int array instead of short array to match MqttSubAckPayload constructor
        int[] qosArray = new int[grantedQosLevels.size()];
        for (int i = 0; i < grantedQosLevels.size(); i++) {
            qosArray[i] = grantedQosLevels.get(i);
        }
        MqttSubAckPayload subAckPayload = new MqttSubAckPayload(qosArray);
        MqttSubAckMessage subAck = new MqttSubAckMessage(subAckFixedHeader, subAckVariableHeader, subAckPayload);
        ctx.writeAndFlush(subAck);

        for (MqttTopicSubscription subscription : topicSubscriptions) {
            String topicFilter = subscription.topicName();
            List<Message> retainMessages = config.getMessageStore().getRetainMessage(topicFilter);
            for (Message retainMessage : retainMessages) {
                // 使用消息分发器发送保留消息
                config.getMessageDispatcher().sendRetainMessage(ctx, clientId, retainMessage);
            }
        }
    }

    public void processUnsubscribe(ChannelHandlerContext ctx, MqttUnsubscribeMessage msg) {
        String clientId = ctx.channel().attr(CLIENT_ID_ATTR).get();
        if (clientId == null) {
            return;
        }

        List<String> topicFilters = msg.payload().topics();
        for (String topicFilter : topicFilters) {
            config.getSessionManager().removeSubscribe(topicFilter, clientId);
            LogTool.debug("Client {} unsubscribed from topic {}", clientId, topicFilter);
        }

        MqttUnsubAckMessage unsubAck = MqttMessageBuilders.unsubAck()
                .packetId(msg.variableHeader().messageId())
                .build();
        ctx.writeAndFlush(unsubAck);
    }

    public void processPingReq(ChannelHandlerContext ctx) {
        MqttFixedHeader pingRespFixedHeader = new MqttFixedHeader(MqttMessageType.PINGRESP, false, MqttQoS.AT_MOST_ONCE, false, 0);
        MqttMessage pingResp = new MqttMessage(pingRespFixedHeader);
        ctx.writeAndFlush(pingResp);
    }

    public void processDisconnect(ChannelHandlerContext ctx) {
        String clientId = ctx.channel().attr(CLIENT_ID_ATTR).get();
        if (clientId != null) {
            config.getChannelManager().removeChannel(clientId);
            config.getSessionManager().remove(clientId);
            LogTool.info("Client {} disconnected", clientId);
        }
        ctx.close();
    }

    private void sendConnAck(ChannelHandlerContext ctx, MqttConnectReturnCode returnCode) {
        MqttConnAckMessage connAck = MqttMessageBuilders.connAck()
                .returnCode(returnCode)
                .sessionPresent(false)
                .build();
        ctx.writeAndFlush(connAck);
        ctx.close();
    }

    public void sendWillMessage(Message willMessage) {
        if (willMessage == null || willMessage.getTopic() == null) {
            return;
        }

        List<Subscribe> subscribers = config.getSessionManager().searchSubscribe(willMessage.getTopic());
        for (Subscribe subscribe : subscribers) {
            String subClientId = subscribe.getClientId();
            ChannelHandlerContext subCtx = config.getChannelManager().getChannel(subClientId);
            if (subCtx == null || !subCtx.channel().isActive()) {
                continue;
            }

            int subQos = subscribe.getMqttQoS();
            MqttQoS finalQos = MqttQoS.valueOf(willMessage.getQos());
            if (finalQos.value() > subQos) {
                finalQos = MqttQoS.valueOf(subQos);
            }

            publishMessage(subCtx, subClientId, willMessage.getTopic(), willMessage.getPayload(),
                    finalQos, willMessage.isRetain(), config.getSerializer());
        }
    }

    public boolean publishMessage(ChannelHandlerContext ctx, String topic, Object payload,
                                  MqttQoS qos, boolean retain, IMessageSerializer serializer) {
        return publishMessage(ctx, null, topic, payload, qos, retain, serializer);
    }

    public boolean publishMessage(ChannelHandlerContext ctx, String clientId, String topic, Object payload,
                                  MqttQoS qos, boolean retain, IMessageSerializer serializer) {
        try {
            byte[] payloadBytes;
            if (payload instanceof byte[]) {
                payloadBytes = (byte[]) payload;
            } else {
                payloadBytes = serializer.serialize(payload);
            }

            int messageId = -1;
            if (qos == MqttQoS.AT_LEAST_ONCE || qos == MqttQoS.EXACTLY_ONCE) {
                if (clientId == null) {
                    clientId = ctx.channel().attr(CLIENT_ID_ATTR).get();
                }
                if (clientId != null) {
                    messageId = config.getSessionManager().getPacketId(clientId);
                }
            }

            MqttFixedHeader publishFixedHeader = new MqttFixedHeader(MqttMessageType.PUBLISH, false, qos, retain, 0);
            MqttPublishVariableHeader publishVariableHeader = new MqttPublishVariableHeader(topic, messageId);
            MqttPublishMessage publishMessage = new MqttPublishMessage(publishFixedHeader, publishVariableHeader, io.netty.buffer.Unpooled.copiedBuffer(payloadBytes));

            if (qos == MqttQoS.AT_LEAST_ONCE && messageId > 0 && clientId != null) {
                MqttPendingPublish pendingPublish = new MqttPendingPublish(publishMessage, qos);
                config.getSessionManager().addPendingPublish(clientId, messageId, pendingPublish);
                pendingPublish.startPublishRetransmissionTimer(scheduledExecutor, ctx);
            } else if (qos == MqttQoS.EXACTLY_ONCE && messageId > 0 && clientId != null) {
                MqttPendingPublish pendingPublish = new MqttPendingPublish(publishMessage, qos);
                config.getSessionManager().addPendingPublish(clientId, messageId, pendingPublish);
                pendingPublish.startPublishRetransmissionTimer(scheduledExecutor, ctx);
            }

            ctx.writeAndFlush(publishMessage);
            return true;
        } catch (Exception e) {
            LogTool.error("Error publishing message", e);
            return false;
        }
    }
}

